import numpy as np
from scipy.spatial.distance import cdist
import time
from concurrent.futures import ThreadPoolExecutor

# 1. 生成向量数据
def generate_vector_database(num_vectors=100000, dimensions=500):
    """生成模拟向量数据库"""
    np.random.seed(42)
    return np.random.rand(num_vectors, dimensions)

vector_database = generate_vector_database(num_vectors=1000000, dimensions=500)
query_vector = np.random.rand(1, 500)  # 模拟一个查询向量

# 2. 基础暴力搜索（标准化）
def brute_force_search(query, database, top_k=5):
    """基础暴力搜索实现"""
    distances = cdist(query, database, metric='euclidean').flatten()
    top_k_indices = np.argsort(distances)[:top_k]
    top_k_distances = distances[top_k_indices]
    return top_k_indices, top_k_distances

# 3. 优化 1：矩阵运算优化
def optimized_brute_force_new(query, database, top_k=5):
    # 更高效的实现
    query_sq = np.sum(query**2, axis=1)  # 形状: (1,)
    database_sq = np.sum(database**2, axis=1)  # 形状: (n,)
    dot_product = np.dot(database, query.T).flatten()  # 形状: (n,)
    
    distances = np.sqrt(query_sq + database_sq - 2 * dot_product)
    
    top_k_indices = np.argsort(distances)[:top_k]
    top_k_distances = distances[top_k_indices]
    return top_k_indices, top_k_distances

def optimized_brute_force(query, database, top_k=5):
    """通过矩阵运算优化暴力搜索"""
    query_norm = np.linalg.norm(query, axis=1, keepdims=True) ** 2
    database_norm = np.linalg.norm(database, axis=1, keepdims=True) ** 2
    distances = np.sqrt(query_norm + database_norm.T - 2 * np.dot(query, database.T)).flatten()
    
    top_k_indices = np.argsort(distances)[:top_k]
    top_k_distances = distances[top_k_indices]
    return top_k_indices, top_k_distances

# 4. 优化 2：并行化计算
def parallel_brute_force(query, database, top_k=5, num_threads=4):
    """通过多线程加速暴力搜索"""
    def compute_distances(chunk):
        return cdist(query, chunk, metric='euclidean').flatten()
    
    # 将数据库划分为块
    chunk_size = len(database) // num_threads
    chunks = [database[i * chunk_size:(i + 1) * chunk_size] for i in range(num_threads)]
    
    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        results = list(executor.map(compute_distances, chunks))
    
    distances = np.concatenate(results)
    top_k_indices = np.argsort(distances)[:top_k]
    top_k_distances = distances[top_k_indices]
    return top_k_indices, top_k_distances

# 5. 运行与性能对比
start_time = time.time()
indices_base, distances_base = brute_force_search(query_vector, vector_database, top_k=5)
time_base = time.time() - start_time

start_time = time.time()
indices_opt, distances_opt = optimized_brute_force_new(query_vector, vector_database, top_k=5)
time_opt = time.time() - start_time

start_time = time.time()
indices_par, distances_par = parallel_brute_force(query_vector, vector_database, top_k=5)
time_par = time.time() - start_time

# 6. 输出结果
print(f"基础暴力搜索 - 时间: {time_base:.4f}s, 索引: {indices_base}, 距离: {distances_base}")
print(f"矩阵优化搜索 - 时间: {time_opt:.4f}s, 索引: {indices_opt}, 距离: {distances_opt}")
print(f"并行暴力搜索 - 时间: {time_par:.4f}s, 索引: {indices_par}, 距离: {distances_par}")